from __future__ import print_function, division
"""
load image/label/box from annotation file
"""

import os, sys, random
import os.path as osp
import numpy as np
from PIL import Image

import torch
from torch.utils.data import Dataset
import torchvision as tv

from encoder import DataEncoder

class ListDataset(Dataset):
    """
    The list file is like:
        img.jpg xmin ymin xmax ymax label xmin ymin xmax ymax label
    """
    def __init__(self, root, list_file, train, transform, input_size):
        """
        initialization
        Input:  root:   (str) directory to images
                list_file (str) path to index file
                train:  (boolean) train or def test
                transform:  ([transforms]) image transforms
                input_size:     (int) model input size
        """

        # output encoded bbox coordinates
        self._useEncode = True

        self._root  = rrot
        self._list_file = list_file
        self._train = train
        self._transform = transform
        self._input_size= input_size

        # define the meta info list
        self._fname_list= []
        self._bbox_list = []
        self._label_list= []
        # define the encoder
        self._encoder   = DataEncoder

    def _read_list_file(self):
        """
        parse the list file to get each image's info
        """
        assert osp.exists(self._list_file), "list_file: {} not exists".format(self._list_file)
        with open(self._list_file) as fid:
            line_list   = fid.readlines(fid)
            self._num_samples   = len(line_list)

        # parse each line
        for line in line_list:
            _i_split    = line.strip().split()
            # get the image name
            _i_name     = _i_split[0]
            self._fname_list.append(_i_name)

            _i_num_bbox     = int((len(_i_split) - 1) / 5)
            _i_bbox_list    = []
            _i_label_list   = []
            for j in range(_i_num_bbox):
                _j_xmin     = _i_split[1 + 5*i]
                _j_xmax     = _i_split[2 + 5*i]
                _j_ymin     = _i_split[3 + 5*i]
                _j_ymax     = _i_split[4 + 5*i]
                _j_class    = _i_split[5 + 5*i]
                # append to current image's info
                _j_bbox     = [float(_j_xmin), float(_j_ymin), float(_j_xmax), float(_j_ymax)]
                _i_bbox_list.append(_j_bbox)
                _i_label_list.append(int(_j_class))
            self._bbox_list.append(_i_bbox_list)
            self._label_list.append(_i_label_list)

    def __getitem__(self, index):
        """
        load image and corresponding label info
        Input:
            index:  (int) image index [0->...]
        Returns:
            img:    image array
            loc_target_list:    (list of array) location target list
            cls_target_list:    (list of array) classification target list
        """
        # load image and bbox
        fname   = self._fname_list[index]
        fpath   = osp.join(self._root, fname)
        with open(fpath) as fid:
            _img= Image.open(fid)
            if _img.mode != "RGB":
                _img    = _img.convert("RGB")
        _bbox_list  = self._bbox_list[index]
        _label_list = self._label_list[index]
        _size       = self._input_size

        # TODO, data augmentation

    def collate_fn(self, batch):

